from llmtcg.agents.testers.os.linux import LinuxSyscallTesterAgent
# from llmtcg.agents.coders.c import CCoderAgent, CodingFailed
from llmtcg.agents.coders.openssl_c import CCoderAgent, CodingFailed
from llmtcg.chat.models.claude import ClaudeChatModel
from llmtcg.chat.models.local import LocalOpenaiChatModel
from llmtcg.chat.models.deepseek import DeepSeekChatModel
from llmtcg.chat.models.siliconflow import SiliconFlowChatModel
from llmtcg.chat.models.groq import GroqChatModel
from llmtcg.chat.models.mistral import MistralChatModel
from llmtcg.chat.models.openai import OpenaiChatModel
from llmtcg.chat.models.yi import YiChatModel
from llmtcg.utils.profiler import Profiler
# from llmtcg.utils.mansplitter import ManSplitter
from llmtcg.utils.openssl_mansplitter import OpenSSLSplitter
from llmtcg.utils.openssl_store import OpenSSLStore

from loguru import logger
from subprocess import run, PIPE
from argparse import ArgumentParser, Namespace
from typing import Optional, Tuple, List
from os import path, makedirs, environ, walk, listdir, remove, system
from sys import exit

from concurrent.futures import as_completed, ThreadPoolExecutor
import multiprocessing
import atexit
import json
import shutil

profiler = Profiler()

tester_num = 32 # 测试设计者线程数
testers = []

coder_num = 128 # 测试执行者线程数
coders = []

ignored_syscalls = []
# unsupportted parameter options: dict[syscall: [unsupportted flags]]
unsupportted_syscalls = {}
openssl_store = OpenSSLStore("gen_funs.json")

_shared_counter = multiprocessing.Manager().Queue()
_shared_total = 0


def parse_args() -> Namespace:
    parser = ArgumentParser(
        "python3 llmtcg.apps.linux_syscall_tcg",
        description="generating test cases for linux kernel syscalls"
    )
    parser.add_argument(
        "--certs", type=str, required=True,
        help="The directory containing the collection of certificates generated by gen_certs.sh"
    )
    parser.add_argument(
        "--syscalls", type=str,
        help="list of syscall names, separated by ',', e.g. 'open,read'. " +
             "By default all syscalls are generated"
    )
    parser.add_argument(
        "--resume", action="store_true",
        help="generate test code for previously failed scenes"
    )
    parser.add_argument(
        "--use-manpage", action="store_true",
        help="use linux manpage to grab syscall documents"
    )
    parser.add_argument(
        "--use-coding-history", action="store_true",
        help="use coding history to improve code generation performance"
    )
    parser.add_argument(
        "--jobs", metavar='N', type=int, default=coder_num,
        help="using multiple works to generate test case in parallel"
    )
    parser.add_argument(
        "-o", "--output", type=str, required=True,
        help="the directory to dump the generated test cases"
    )
    return parser.parse_args()


def initialize(args: Namespace) -> None:
    index = 0
    while index < tester_num:
        new_tester_model = SiliconFlowChatModel(model=SiliconFlowChatModel.Model.SILICONFLOW_DS_V3_PRO)
        new_tester_model.api_key = environ.get(f"SILICONFLOW_API_KEY_{index % 16}", "")
        new_tester_model.profiler = profiler
        new_tester = LinuxSyscallTesterAgent(new_tester_model)
        testers.append(new_tester)
        index += 1
    
    index = 0
    while index < args.jobs:
        new_coder_model = SiliconFlowChatModel(model=SiliconFlowChatModel.Model.SILICONFLOW_DS_V3_PRO)
        new_coder_model.api_key = environ.get(f"SILICONFLOW_API_KEY_{index % 16}", "")
        new_coder_model.profiler = profiler

        new_coder = CCoderAgent(new_coder_model)
        new_coder.compiler_args.append("-Werror")
        coders.append(new_coder)
        index += 1
    
    for coder in coders:
        coder.with_history = args.use_coding_history
    logger.info(f"use coding history: {args.use_coding_history}")

def get_manpage(syscall: str) -> Optional[str]:
    proc = run(f"man 3 {syscall}", shell=True, stdin=PIPE, stdout=PIPE, stderr=PIPE)
    if proc.returncode != 0:
        logger.warning(f"failed to get manpage of syscall {syscall}")
        return None

    content = proc.stdout.decode()
    return content


def generate_test_code(scene: str, target_c_path: str, i: int = 0) -> None:
    index = _shared_counter.qsize()
    coder = coders[i % len(coders)]
    try:
        logger.debug(f"generating test code for {target_c_path}")
        prompt = f"""

        Generate a c testing program that implements the following test scene: {scene}.
        The testing program should:

        - Directly use the assert() function to give the assertions and check error code, instead of using printf() function.
        - Never try to invoke shell commands through `system` or simular functions.
        - Must assume that the user give no input from stdin.
        - Before using OPENSSL API, the corresponding header file must be included. For example, to use functions related to protocol, include <openssl/ssl.h>; to use functions related to error handling, include <openssl/err.h>. It depends on the called function, but remember to include the correct OpenSSL header files.
        - Before using OPENSSL API, make sure it exists. Only use OPENSSL API that can be found in the official documentation, manual or header files of OPENSSL.
        - When invoking OpenSSL API, work really hard to invoke them in accordance with the prototypes specified in OpenSSL's header files.
        - Must access the members of OPENSSL structure only through the get function provided by OPENSSL.
        - Before introducing memory operations or passing NULL pointers to the OpenSSL API, please carefully consider if this might cause segmentation faults or program crashes. If so, do not do it.

        """

        if openssl_store.check_cert_key(scene):
            prompt += f"""
            Remember the following certificate or key directory:
            - "{openssl_store.get_certs_path()}/algorithms/rsa_2048" : the directory for storing RSA (2048 bits) certificates and keys
            - "{openssl_store.get_certs_path()}/algorithms/rsa_4096" : the directory for storing RSA (4096 bits) certificates and keys
            - "{openssl_store.get_certs_path()}/algorithms/rsa_pss_2048" : the directory for storing RSA-PSS (2048 bits) certificates and keys
            - "{openssl_store.get_certs_path()}/algorithms/rsa_pss_4096" : the directory for storing RSA-PSS (4096 bits) certificates and keys
            - "{openssl_store.get_certs_path()}/algorithms/ed25519" : the directory for storing ED25519 certificates and keys
            - "{openssl_store.get_certs_path()}/algorithms/ed448" : the directory for storing ED448 certificates and keys
            - "{openssl_store.get_certs_path()}/algorithms/ec_p256" : the directory for storing EC-P256 certificates and keys
            - "{openssl_store.get_certs_path()}/algorithms/ec_p384" : the directory for storing EC-P384 certificates and keys
            These directories store the following certificates or key files: 
            - "root_ca.crt" : root CA certificate
            - "root_ca.key" : root CA certificate's private key
            - "int_ca.crt" : intermediate CA certificate
            - "int_ca.key" : intermediate CA certificate's private key
            - "server.crt" : server certificate
            - "server.key" : server certificate's private key
            - "client.crt" : client certificate
            - "client.key" : client certificate's private key
            - "*.csr" : Certificate Signaing request (CSR) corresponding to the certificate
            - "*.cnf" : Configuration corresponding to the certificate
            - "*.der" : DER (Distinguished Encoding Rules) corresponding to the certificate
            - "*.p7b" : PKCS#7 format corresponding to the certificate
            - "*.pfx" : PKCS#12 format corresponding to the certificate and key
            - "*.crl" : Certificate Revocation List containing the corresponding certificate
            """

        program = coder.get_program(prompt)
        logger.info(f"({index}/{_shared_total}) succeeded to generate {target_c_path}")

        with open(target_c_path, "w+") as fprog:
            fprog.write(program)

        with open(target_c_path + ".log", "w+") as flog:
            summary_json = coder.last_context.summary_json()
            flog.write(json.dumps(summary_json))

        _shared_counter.put((target_c_path, True))

    except CodingFailed as ex:
        logger.error(f"({index}/{_shared_total}) failed to generate {target_c_path}: {ex}")
        with open(f"{target_c_path}.err", "w+") as ferr:
            ferr.write(f"error: {ex}")
            ferr.write(coder.last_context.summary())

        _shared_counter.put((target_c_path, False))

def generate_scene_prompts(syscall: str, manpage: str) -> list[str]:
    scene_prompts = []
    splitter = OpenSSLSplitter(manpage)
    splitter.parse()

    for section_title, contents in splitter.get_parsed_data().items():
        for content in contents:
            scene_prompt = f"""
            Generate a full list of test scenes of the openssl api based on the manpage chapter shown below.\n

            ```
            {content}
            ```
            """
        
            # 接口测试/组合测试/参数测试/返回值测试/规格测试
            scene_prompt_general = scene_prompt + f"""
            You need to create as many test scenes as possible, and work really hard to explore every potentially behavior of OpenSSL API. To cover normal, corner, and invalid usage, you need to satisfy the following points: 
            - Cover every function point mentioned in the above manpage chapter. 
            - Consider every combinational invocation with one or more functions mentioned in manpage, especially "SYNOPSIS" and "SEE ALSO" sections. Work really hard to explore every potentially behavior of OPENSSL API under combinational scenarios, with emphasis on call order dependencies (e.g., functions that must be called after initialization) and state transitions (e.g., how a function's behavior changes when the context state is modified by another function). 
            - For each parameter, make sure that for normal and corner assignments (e.g., null pointers, boundary values, empty buffers, special characters), OpenSSL API works well and the functionality is consistent with manpage; and for invalid assignment, the behavior is consistent with manpage (e.g., returned error code is correct). 
            - Cover every invalid usage mentioned in manpage is addressed as expected. Do not include any test scenarios that are explicitly documented to cause undefined behavior (e.g., dereferencing a null pointer when the manual explicitly states "undefined behavior") or are known to crash the program. Focus on scenarios where the API is expected to handle errors gracefully (e.g., return an error code or set an error state). 
            - The return values of OpenSSL API are strictly consistent with manpage, with particular attention to the "RETURN VALUES" section. For verification: check return codes, output data validity, and error codes as needed. 
            - Cover all the specifications, options and modes mentioned in manpage, including algorithm-specific constraints. 
            - Each test scene must include: test objective (the API behavior/function to verify), preconditions (e.g., initialization state, dependent objects), operation steps (function call order, parameter assignment method), expected results (return values, output data, error codes and messages), and notes (e.g., relationship with other scenes). 
            - Classify all scenes by functional module (e.g., symmetric encryption), error type (e.g., invalid parameters), and call mode (e.g., sequential calls) to avoid duplication. 
            """
            scene_prompts.append(scene_prompt_general)

            # 并发测试
            if openssl_store.check_thread_safety(syscall, content):
                scene_prompt_concurrent = scene_prompt + f"""
                You need to generate as many test scenes related to concurrency and concurrent invocation as possible, you need to satisify the following points:
                - Cover every function point mentioned in the above manpage chapter, 
                - Cover normal, corner, and invalid usage, 
                - Work really hard to explore every potentially concurrent behavior of this openssl api, 
                """
                scene_prompts.append(scene_prompt_concurrent)

    return scene_prompts

def filter_dup_functions(manpage: str) -> bool:
    generated_funs = openssl_store.get_generated_functions()
    function_name = manpage.split('(', 1)[0]
    logger.info(f"fun name: {function_name}")
    if function_name in generated_funs:
        return True
    else:
        generated_funs.append(function_name)
        return False

def generate_syscall(syscall: str, use_manpage: bool, target_dir: str) -> None:
    if path.exists(target_dir):
        logger.warning(f"Skip {syscall} since it already exists")
        return
    
    manpage = get_manpage(syscall)
    if manpage is None and use_manpage:
        logger.error(f"cannot find manpage for {syscall}, it could not be a valid syscall")
        return

    if (filter_dup_functions(manpage)):
        logger.info(f"duplicate function")
        return
    
    makedirs(target_dir, exist_ok=True)

    scene_prompts = generate_scene_prompts(syscall, manpage)
    logger.info(f"{len(scene_prompts)} scene_prompts for openssl api {syscall}")

    scenes = []
    executor = ThreadPoolExecutor(max_workers=tester_num)
    futures = []
    try:
        for i, scene_prompt in enumerate(scene_prompts):
            scene_prompt += "\n Make sure each test scene meets the specified format requirements.\n"
            futures.append(
                executor.submit(testers[i % tester_num].generate_test_scene, scene_prompt))
        logger.info(f"submitting {len(scene_prompts)} scene prompts for {syscall}")

        completed_cnt = 0
        for future in as_completed(futures):
            try:
                scenes += future.result()
                completed_cnt += 1
                logger.info(f"Scene {completed_cnt} generation success.")
            except Exception as e:
                logger.error(f"Scene generation for prompt because: {e}")
        logger.info(f"{len(scenes)} test scenes generated for openssl api {syscall}")

    finally:
        executor.shutdown(wait=True)

    def unsupportted_filter(scenes: List[str]):
        if syscall not in unsupportted_syscalls:
            return
        wait_list = []
        for scene in scenes:
            for unsupport_flag in unsupportted_syscalls[syscall]:
                if unsupport_flag in scene:
                    wait_list.append(scene)
                    break
        scenes[:] = [scene for scene in scenes if scene not in wait_list]
    
    unsupportted_filter(scenes)
    logger.info(f"{len(scenes)} test scenes left for openssl api {syscall}")
    
    executor = ThreadPoolExecutor(max_workers=coder_num)
    futures = []
    try:
        for i, scene in enumerate(scenes):
            with open(path.join(target_dir, f"test_case_{i}.txt"), "w+") as fdesc:
                fdesc.write(scene)

            c_program_path = path.join(target_dir, f"test_case_{i}.c")
            futures.append(executor.submit(generate_test_code, scene, c_program_path, i))
        
        for future in as_completed(futures):
            try:
                future.result(timeout=180)
            except Exception as e:
                logger.info(f"Task failed because of :{e}\n")
    finally:
        executor.shutdown(wait=True)

    try:
        run("find /tmp | sudo -S xargs rm -rf", 
            input = environ.get('root_pwd')+"\n",
            shell=True, check=True, timeout=30, capture_output=True, text=True)
    except Exception as e:
        logger.info(f"remove remprary files failed because: {e}")


def resume_generate(args: Namespace) -> None:
    global _shared_total
    global coder_num

    try:
        run("find /tmp | sudo -S xargs rm -rf", 
            input = environ.get('root_pwd')+"\n",
            shell=True, check=True, timeout=30, capture_output=True, text=True)
    except Exception as e:
        logger.info(f"remove remprary files failed because: {e}")

    target_dir = args.output
    coder_num = args.jobs

    jobs: List[Tuple[str, str]] = []

    proc = run(f"find {target_dir} | grep .txt$", shell=True, stdout=PIPE, stderr=PIPE)
    scene_paths = proc.stdout.decode().splitlines()
    logger.info(f"find {len(scene_paths)} test scenes in total")

    def get_content(path: str):
        with open(path, "r") as reader:
            content = reader.read()
        return content

    for scene_path in scene_paths:
        case_path = scene_path[:-4] + ".c"
        if path.isfile(case_path):
            # logger.debug(f"test case of {scene_path} has already been generated")
            continue
        elif path.isfile(f"{case_path}.err"):
            content = get_content(f"{case_path}.err")
            if "chat API failed" not in content:
                # logger.debug(f"ignore {scene_path} as it has encountered an error before")
                continue
            try:
                remove(f"{case_path}.err")
            except Exception as e:
                pass
                
        scene = get_content(scene_path)
        jobs.append((scene, case_path))

    _shared_total = len(jobs)

    logger.info(f"using {coder_num} workers to generate test cases for {_shared_total} test scenes")
    executor = ThreadPoolExecutor(max_workers=coder_num)
    futures = []
    # if len(jobs) > 720:
    #     jobs = jobs[:720]
    try:
        for i, (scene, case_path) in enumerate(jobs):
            futures.append(executor.submit(generate_test_code, scene, case_path, i))
        
        for future in as_completed(futures):
            try:
                future.result(timeout=180)
            except Exception as e:
                logger.info(f"Task failed because of :{e}\n")
    finally:
        executor.shutdown(wait=True)


def generate_syscalls(args: Namespace) -> None:

    if args.syscalls is None:
        syscall_names = []
        syscall_names = [syscall for syscall in syscall_names if syscall not in ignored_syscalls]
    else:
        syscall_names = sorted(s.strip() for s in args.syscalls.split(","))

    output_dir = args.output
    logger.info(f"generating test cases for {syscall_names}")
    logger.info(f"there are {len(syscall_names)} syscalls in total")
        
    for syscall in syscall_names:
        syscall_target_dir = path.join(output_dir, syscall)
        generate_syscall(syscall, args.use_manpage, syscall_target_dir)
        

def print_stats() -> None:
    profiler.print_stats()

def delete_c_files_if_unsupportted(root_dir):

    def contains_unsupportted_flags(file_path):
        with open(file_path, 'r', encoding='utf-8') as file:
            content = file.read()
        ignored_syscalls_set = set(ignored_syscalls)
        if any(syscall in content for syscall in ignored_syscalls_set):
            return True
        
        for syscall, flags in unsupportted_syscalls.items():
            if syscall in content and any(flag in content for flag in flags):
                return True

        return False
    
    for syscall_folder in listdir(root_dir):
        syscall_path = path.join(root_dir, syscall_folder)

        if path.isdir(syscall_path) and syscall_folder in ignored_syscalls:
            shutil.rmtree(syscall_path)
            continue
        
        cnt = 0
        for dirpath, _, filenames in walk(syscall_path):
            for filename in filenames:
                if filename.endswith('.c'):
                    file_path = path.join(dirpath, filename)
                    if contains_unsupportted_flags(file_path):
                        remove(file_path)
                        cnt += 1

                        scene_file_path = file_path.replace('.c', '.txt')
                        if path.exists(scene_file_path):
                            remove(scene_file_path)

        if cnt > 0:
            logger.info(f"remove {cnt} testcases for {syscall_folder} due to unsupportted flags")


if __name__ == "__main__":
    args = parse_args()
    is_valid = all([curr_coder.chat_model.validate() for curr_coder in coders])
    if not is_valid:
        logger.critical(f"model is not properly configured")
        exit(1)

    atexit.register(print_stats)
    initialize(args)
    openssl_store.set_certs_path(args.certs)

    if args.resume:
        resume_generate(args)
    else:
        generate_syscalls(args)

    delete_c_files_if_unsupportted(args.output)
    # 更新 OPENSSL 已处理的接口列表
    openssl_store.write_json_gen_funs("gen_funs.json")
